{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_core.messages import HumanMessage, SystemMessage\n",
    "from langgraph.graph import MessagesState\n",
    "from langgraph.graph import START, END, StateGraph\n",
    "from IPython.display import Image, display\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "import json\n",
    "\n",
    "\n",
    "# OPENAI_API_KEY environment variable must be set\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "\n",
    "\n",
    "# Defining Agent's node\n",
    "##################################################################################\n",
    "\n",
    "# System message\n",
    "chatbot_system_message = SystemMessage(content=(\"\"\"\n",
    "You are a helpful and knowledgeable chatbot assistant. \n",
    "Your goal is to provide clear and accurate answers to user questions based on the information they provide. \n",
    "Stay focused, concise, and ensure your responses are relevant to the context of the conversation. \n",
    "If you don’t have enough information, ask for clarification.”\n",
    "\"\"\"))\n",
    "\n",
    "# Node\n",
    "def chatbot(state: MessagesState):\n",
    "   return {\"messages\": [llm.invoke([chatbot_system_message] + state[\"messages\"])]}\n",
    "\n",
    "\n",
    "# Defining Graph\n",
    "##################################################################################\n",
    "\n",
    "builder = StateGraph(MessagesState)\n",
    "builder.add_node(\"chatbot\", chatbot)\n",
    "\n",
    "\n",
    "# Define edges: these determine how the control flow moves\n",
    "builder.add_edge(START, \"chatbot\")\n",
    "builder.add_edge(\"chatbot\", END)\n",
    "\n",
    "memory = MemorySaver()\n",
    "chatbot_graph = builder.compile(checkpointer=memory)\n",
    "\n",
    "# Show\n",
    "display(Image(chatbot_graph.get_graph(xray=True).draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thread_id = \"2\"\n",
    "\n",
    "config = {\"configurable\": {\"thread_id\": thread_id}}\n",
    "\n",
    "chatbot_graph.invoke({\"messages\": [HumanMessage(content=\"Hi, I’m working on a Python project, and I’m stuck with handling API responses.\")]}, config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chatbot_graph.invoke({\"messages\": [HumanMessage(content=\"Sorry what was my previous question?\")]}, config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chatbot_graph.invoke({\"messages\": [HumanMessage(content=\"Ahh, yeah right! So I’m mostly struggling with parsing JSON responses. Sometimes the structure isn’t what I expect, and it breaks my code.\")]}, config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chatbot_graph.invoke({\"messages\": [HumanMessage(content=\"Got it! That helps a lot. What would be the best way to handle deeply nested JSON data when I only need a few specific values?\")]}, config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Defining your own schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_core.messages import HumanMessage, SystemMessage\n",
    "from langgraph.graph import START, END, StateGraph\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from IPython.display import Image, display\n",
    "\n",
    "# OPENAI_API_KEY environment variable must be set\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "\n",
    "# Defining Schema\n",
    "##################################################################################\n",
    "class State(MessagesState):\n",
    "    question: str\n",
    "    answer: str\n",
    "\n",
    "\n",
    "# Defining Agent's node\n",
    "##################################################################################\n",
    "\n",
    "# System message\n",
    "chatbot_system_message = SystemMessage(content=(\"\"\"\n",
    "You are a helpful and knowledgeable chatbot assistant. \n",
    "Your goal is to provide clear and accurate answers to user questions based on the information they provide. \n",
    "Stay focused, concise, and ensure your responses are relevant to the context of the conversation. \n",
    "If you don’t have enough information, ask for clarification.”\n",
    "\"\"\"))\n",
    "\n",
    "# Node\n",
    "def chatbot(state: State) -> State:\n",
    "    question = HumanMessage(content=state.get(\"question\", \"\"))\n",
    "    response = llm.invoke([chatbot_system_message] + state[\"messages\"] + [question]);\n",
    "\n",
    "    return State(\n",
    "        messages = [question, response],\n",
    "        question = state.get(\"question\", None),\n",
    "        answer = response.content\n",
    "    )\n",
    "\n",
    "\n",
    "# Defining Graph\n",
    "##################################################################################\n",
    "\n",
    "builder = StateGraph(State)\n",
    "builder.add_node(\"chatbot\", chatbot)\n",
    "\n",
    "\n",
    "# Define edges: these determine how the control flow moves\n",
    "builder.add_edge(START, \"chatbot\")\n",
    "builder.add_edge(\"chatbot\", END)\n",
    "\n",
    "memory = MemorySaver()\n",
    "chatbot_graph = builder.compile(checkpointer=memory)\n",
    "\n",
    "# Show\n",
    "display(Image(chatbot_graph.get_graph(xray=True).draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thread_id = \"2\"\n",
    "\n",
    "config = {\"configurable\": {\"thread_id\": thread_id}}\n",
    "\n",
    "chatbot_graph.invoke(State(question=\"Hi, I’m working on a Python project, and I’m stuck with handling API responses.\"), config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chatbot_graph.invoke(State(question=\"Sorry what was my previous question?\"), config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chatbot_graph.invoke(State(question=\"Ahh, yeah right! So I’m mostly struggling with parsing JSON responses. Sometimes the structure isn’t what I expect, and it breaks my code.\"), config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chatbot_graph.invoke(State(question=\"Got it! That helps a lot. What would be the best way to handle deeply nested JSON data when I only need a few specific values?\"), config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Message Shrinking"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Limit number of messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_core.messages import HumanMessage, SystemMessage\n",
    "from langgraph.graph import START, END, StateGraph\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from IPython.display import Image, display\n",
    "\n",
    "# OPENAI_API_KEY environment variable must be set\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "\n",
    "# Defining Schema\n",
    "##################################################################################\n",
    "class State(MessagesState):\n",
    "    question: str\n",
    "    answer: str\n",
    "\n",
    "\n",
    "# Defining Agent's node\n",
    "##################################################################################\n",
    "\n",
    "# System message\n",
    "chatbot_system_message = SystemMessage(content=(\"\"\"\n",
    "You are a helpful and knowledgeable chatbot assistant. \n",
    "Your goal is to provide clear and accurate answers to user questions based on the information they provide. \n",
    "Stay focused, concise, and ensure your responses are relevant to the context of the conversation. \n",
    "If you don’t have enough information, ask for clarification.”\n",
    "\"\"\"))\n",
    "\n",
    "# Node\n",
    "def chatbot(state: State) -> State:\n",
    "    question = HumanMessage(content=state.get(\"question\", \"\"))\n",
    "\n",
    "    trimed_messages = (state[\"messages\"] + [question])[-4:]\n",
    "    print(f\"Messages sent to LLM: {len(trimed_messages)}\")\n",
    "\n",
    "    response = llm.invoke([chatbot_system_message] + trimed_messages);\n",
    "\n",
    "    return State(\n",
    "        messages = [question, response],\n",
    "        question = state.get(\"question\", None),\n",
    "        answer = response.content\n",
    "    )\n",
    "\n",
    "\n",
    "# Defining Graph\n",
    "##################################################################################\n",
    "\n",
    "builder = StateGraph(State)\n",
    "builder.add_node(\"chatbot\", chatbot)\n",
    "\n",
    "\n",
    "# Define edges: these determine how the control flow moves\n",
    "builder.add_edge(START, \"chatbot\")\n",
    "builder.add_edge(\"chatbot\", END)\n",
    "\n",
    "memory = MemorySaver()\n",
    "chatbot_graph = builder.compile(checkpointer=memory)\n",
    "\n",
    "# Show\n",
    "display(Image(chatbot_graph.get_graph(xray=True).draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thread_id = \"1\"\n",
    "\n",
    "config = {\"configurable\": {\"thread_id\": thread_id}}\n",
    "\n",
    "chatbot_graph.invoke(State(question=\"Hi, I’m working on a Python project, and I’m stuck with handling API responses.\"), config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chatbot_graph.invoke(State(question=\"Sorry what was my previous question?\"), config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chatbot_graph.invoke(State(question=\"Ahh, yeah right! So I’m mostly struggling with parsing JSON responses. Sometimes the structure isn’t what I expect, and it breaks my code.\"), config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trim messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_core.messages import HumanMessage, SystemMessage\n",
    "from langgraph.graph import START, END, StateGraph\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from IPython.display import Image, display\n",
    "from langchain_core.messages import trim_messages\n",
    "\n",
    "# OPENAI_API_KEY environment variable must be set\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "\n",
    "# Defining Schema\n",
    "##################################################################################\n",
    "class State(MessagesState):\n",
    "    question: str\n",
    "    answer: str\n",
    "\n",
    "\n",
    "# Defining Agent's node\n",
    "##################################################################################\n",
    "\n",
    "# System message\n",
    "chatbot_system_message = SystemMessage(content=(\"\"\"\n",
    "You are a helpful and knowledgeable chatbot assistant. \n",
    "Your goal is to provide clear and accurate answers to user questions based on the information they provide. \n",
    "Stay focused, concise, and ensure your responses are relevant to the context of the conversation. \n",
    "If you don’t have enough information, ask for clarification.”\n",
    "\"\"\"))\n",
    "\n",
    "# Node\n",
    "def chatbot(state: State) -> State:\n",
    "    question = HumanMessage(content=state.get(\"question\", \"\"))\n",
    "\n",
    "    trimed_messages = trim_messages(\n",
    "            state[\"messages\"] + [question],\n",
    "            max_tokens=200,                # keep 4-5 messages including longer detailed responses\n",
    "            strategy=\"last\",                # keep the most recent\n",
    "            token_counter=llm,              # the model used to count tokens\n",
    "            allow_partial=False             # only include complete messages (preserve message integrity)\n",
    "        )\n",
    "    print(f\"Messages sent to LLM: {len(trimed_messages)}\")\n",
    "\n",
    "    response = llm.invoke([chatbot_system_message] + trimed_messages);\n",
    "\n",
    "    return State(\n",
    "        messages = [question, response],\n",
    "        question = state.get(\"question\", None),\n",
    "        answer = response.content\n",
    "    )@\n",
    "\n",
    "\n",
    "# Defining Graph\n",
    "##################################################################################\n",
    "\n",
    "builder = StateGraph(State)\n",
    "builder.add_node(\"chatbot\", chatbot)\n",
    "\n",
    "\n",
    "# Define edges: these determine how the control flow moves\n",
    "builder.add_edge(START, \"chatbot\")\n",
    "builder.add_edge(\"chatbot\", END)\n",
    "\n",
    "memory = MemorySaver()\n",
    "chatbot_graph = builder.compile(checkpointer=memory)\n",
    "\n",
    "# Show\n",
    "display(Image(chatbot_graph.get_graph(xray=True).draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thread_id = \"3\"\n",
    "\n",
    "config = {\"configurable\": {\"thread_id\": thread_id}}\n",
    "\n",
    "chatbot_graph.invoke(State(question=\"Hi, I’m working on a Python project, and I’m stuck with handling API responses.\"), config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chatbot_graph.invoke(State(question=\"Sorry what was my previous question?\"), config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chatbot_graph.invoke(State(question=\"Ahh, yeah right! So I’m mostly struggling with parsing JSON responses. Sometimes the structure isn’t what I expect, and it breaks my code.\"), config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chatbot_graph.invoke(State(question=\"Got it! That helps a lot. What would be the best way to handle deeply nested JSON data when I only need a few specific values?\"), config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "let's allow partial!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
